import rclpy
from rclpy.node import Node
from sensor_msgs.msg import Image
from std_msgs.msg import String
from cv_bridge import CvBridge
from rclpy.qos import QoSProfile, QoSReliabilityPolicy
import os
import cv2
import numpy as np

def match_direction(image, template_path):
    # 读取模板图片
    template = cv2.imread(template_path)
    if template is None:
        raise ValueError(f"无法读取模板图片: {template_path}")
    
    # 进行模板匹配
    result = cv2.matchTemplate(image, template, cv2.TM_CCOEFF_NORMED)
    
    # 获取最大匹配值
    min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(result)
    
    return max_val  # 返回最大匹配值作为分数

def get_best_template_match(cv_image, templates_dir):
    best_score = float('-inf')
    best_template_num = 0
    
    # 依次匹配四个模板
    for i in range(1, 5):
        template_path = os.path.join(templates_dir, f'template{i}.png')
        score = match_direction(cv_image, template_path)
        if score > best_score:
            best_score = score
            best_template_num = i
    
    return best_template_num

class DirectionMatcherNode(Node):
    def __init__(self):
        super().__init__('direction_matcher_node')

        # 创建best effort QoS
        qos_profile = QoSProfile(
            reliability=QoSReliabilityPolicy.BEST_EFFORT
        )

        # 固定模板路径
        self.templates_dir = "src/direction_detector/B_template"

        # 初始化
        self.bridge = CvBridge()
        self.image_sub = self.create_subscription(
            Image,
            '/rgb_camera/image_raw',
            self.image_callback,
            qos_profile
        )

        self.direction_pub = self.create_publisher(
            String,
            '/direction',
            10
        )

        self.get_logger().info("方向匹配节点已启动。订阅图像话题 /image_raw，发布方向到 /direction")

    def image_callback(self, msg):
        try:
            cv_image = self.bridge.imgmsg_to_cv2(msg, desired_encoding='bgr8')
            
            # 获取最佳匹配的模板序号
            best_template = get_best_template_match(cv_image, self.templates_dir)
            
            # 根据模板序号判断方向：1-2是左，3-4是右
            direction = "left" if best_template <= 2 else "right"

            # 发布方向
            direction_msg = String()
            direction_msg.data = direction
            self.direction_pub.publish(direction_msg)

            self.get_logger().info(f"最佳匹配模板: template{best_template}, 方向: {direction}")

        except Exception as e:
            self.get_logger().error(f"图像处理失败: {e}")


def main(args=None):
    rclpy.init(args=args)
    node = DirectionMatcherNode()
    rclpy.spin(node)
    node.destroy_node()
    rclpy.shutdown()


if __name__ == '__main__':
    main()
